package com.vipboot.framework.core.util.mdc;

import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;

import java.util.Map;

public class MdcTaskDecorator implements TaskDecorator {
    @Override
    public Runnable decorate(Runnable runnable) {
        // 复制当前线程的 MDC 上下文
        final Map<String, String> mdcContext = MDC.getCopyOfContextMap();
        return () -> {
            try {
                // 将 MDC 上下文传递给子线程
                if (mdcContext != null) {
                    MDC.setContextMap(mdcContext);
                }
                // 执行任务
                runnable.run();
            } finally {
                // 清除 MDC 上下文，避免内存泄漏
                MDC.clear();
            }
        };
    }
}
